
#include "ednn_output.h"
#include "ednn_port.h"


static ednn_ret output_forward(struct ednn_layer *layer)
{
    struct ednn_layer *prev = ednn_list_entry(layer->list.prev, 
        struct ednn_layer, list);
    ednn_memcpy(layer->output->pd, prev->output->pd, ednn_tensor_size(prev->output));
    return ednn_ok;
}

static ednn_ret output_build(struct ednn_layer *layer)
{
    struct ednn_layer *prev = ednn_list_entry(layer->list.prev, 
        struct ednn_layer, list);

    if (ednn_tensor_size(prev->output) != ednn_tensor_size(layer->output)) {
        ednn_err("input:%d but need:%d", 
            ednn_tensor_size(prev->output), 
            ednn_tensor_size(layer->output)
        );
        return ednn_err_param;
    }

    return ednn_ok;
}

struct ednn_layer *ednn_output_layer(ednn_shape3d_t shape)
{
    struct ednn_layer *layer = ednn_mem_zalloc(sizeof(struct ednn_layer));
    ednn_assert(layer != EDNN_NULL);

    layer->type    = EDNN_LAYER_OUTPUT;
    layer->forward = output_forward;
    layer->build   = output_build;

    ednn_uint16_t dim[3] = {shape.h, shape.w, shape.c};
    layer->output = ednn_tensor_create(dim, 3);

    return layer;
}